#include "Global_Include.h"
#include "AST_QT.h"

#include "diskpart.h"
#include "includes/CreateBcdBoot.h"
#include "includes/inicpp.hpp"
#include "includes/idFile.h"
#include <boost/filesystem.hpp>

long long getFileSize2(const char* fileName);
std::vector<WCHAR> AstGetUsedLetter();

/// <summary>
/// 复制文件
/// </summary>
/// <param name="Src">源文件</param>
/// <param name="Dst">新文件</param>
/// <returns></returns>
bool AstCopyFile(
	std::wstring Src,
	std::wstring Dst
);
std::wstring AstGetFileExtendName(
	std::wstring filePath
);
void RestorePartition(
	WCHAR CreatedPartition,
	WCHAR ShrinkedPartition
);


/// <summary>
/// 创建PE启动盘
/// </summary>
/// <param name="WimFilePath">PE wim文件路径</param>
/// <param name="SdiFilePath">Sdi文件路径</param>
/// <param name="ImageFilePath">映像文件路径</param>
/// <param name="SelectedIndex">安装的映像索引</param>
/// <param name="SystemReinstallPartition">重装系统驱动器盘符</param>
/// <param name="EfiPartitionLetter">Efi驱动器盘符</param>
/// <param name="ExtendCopyFileList">复制文件</param>
/// <param name="CopyLogToDesktop">是否将日志拷贝到桌面</param>
/// <returns>成功返回true，失败返回false</returns>
bool AST_QT::CreatePePartition(
	std::wstring WimFilePath,
	std::wstring SdiFilePath,

	std::wstring ImageFilePath,
	int SelectedIndex,

	WCHAR SystemReinstallPartition,
	WCHAR EfiPartitionLetter,

	std::vector<CopyFileHandle> ExtendCopyFileList

//	bool CopyLogToDesktop
) {
	AstLog_WriteMessage(
		LogHandle,
		L"[RUNTIME][CREATEPE]开始创建PE分区！"
	);
	AstLog_WriteMessage(
		LogHandle,
		L"[RUNTIME][CREATEPE]开始计算PE分区大小！"
	);
	/* 计算分区大小 */
	unsigned long long TotalPeSpaceByte = 0;

	try {
		//计算大小
		TotalPeSpaceByte += getFileSize2(_wstr2str(WimFilePath).c_str());
		TotalPeSpaceByte += getFileSize2(_wstr2str(SdiFilePath).c_str());
		TotalPeSpaceByte += getFileSize2(_wstr2str(ImageFilePath).c_str());
		for (auto InfoStruct : ExtendCopyFileList) {
			TotalPeSpaceByte += getFileSize2(_wstr2str(InfoStruct.Des).c_str());
		}
	}
	catch (std::wstring ErrInfo) {
		//捕获到异常，报错返回
		AstThrowError_ThrowError(
			ERROR_RUNTIME_CREATEPE_PARTITION_GET_FILE_SIZE_FAIL,
			std::format(L"[RUNTIME][CREATEPE][PARTITION] {}", ErrInfo)
		);
		return false;
	}
	TotalPeSpaceByte = TotalPeSpaceByte;
	unsigned long long TotalPeSpaceMB = round(TotalPeSpaceByte / 1000 / 1000);
	TotalPeSpaceMB = TotalPeSpaceMB + (TotalPeSpaceMB % 24);	//对齐到24MB的整数倍
	AstLog_WriteMessage(
		LogHandle,
		std::format(L"[RUNTIME][CREATEPE][PARTITION]PE分区最终大小:{} MB", TotalPeSpaceMB)
	);

	ui.progressBar_Status->setValue(10);

	/* 创建分区 */
	AstLog_WriteMessage(
		LogHandle,
		L"[RUNTIME][CREATEPE]开始创建PE分区！"
	);

	//接收盘符
	WCHAR PePartitionLetter = NULL;
	WCHAR ShrinkPartitionLetter = NULL;
	//创建分区
	bool bStatus = CreatePeDiskW(&TotalPeSpaceMB, &PePartitionLetter, &ShrinkPartitionLetter);
	//创建失败，报错
	if (!bStatus) {
		AstThrowError_ThrowError(
			ERROR_RUNTIME_CREATEPE_DISKPART_FAIL,
			L"[RUNTIME][CREATEPE][DISKPART]创建PE启动盘失败，详情请见本目录“Log\\DiskpartLog”文件夹下的相关日志。"
		);
		return false;
	}


	ui.progressBar_Status->setValue(30);


	/* 复制文件 */
	AstLog_WriteMessage(
		LogHandle,
		L"[RUNTIME][CREATEPE][COPYFILE]开始复制文件！"
	);

	//创建目录
	CreateDirectory(
		std::format(L"{}:\\WinPE", PePartitionLetter).c_str(),
		NULL
	);

	UINT FileCount = 3 + ExtendCopyFileList.size();		//文件数
	UINT NFileCount = 0;

	std::wstring PePartitionPath = std::format(L"{}:\\", PePartitionLetter);	//PE路径

	AstLog_WriteMessage(
		LogHandle,
		std::format(L"[RUNTIME][CREATEPE][COPYFILE]共有 {} 个文件需要复制", FileCount)
	);

	/* 复制PE文件 */
	//先获取映像后缀
	std::wstring ImageExtendName = AstGetFileExtendName(ImageFilePath);
	bStatus = AstCopyFile(
		ImageFilePath,
		std::format(L"{}Install.{}", PePartitionPath, ImageExtendName)
	);

	try {
		//PE Wim文件
		bStatus = AstCopyFile(
			WimFilePath,
			std::format(L"{}WinPE\\boot.wim", PePartitionPath)
		);
		NFileCount++;
		if (!bStatus) { RestorePartition(PePartitionLetter, ShrinkPartitionLetter); return false; }
		ui.progressBar_Status->setValue(30 + (int)(((NFileCount / FileCount) * 50)));

		//SDI文件
		bStatus = AstCopyFile(
			SdiFilePath,
			std::format(L"{}WinPE\\boot.sdi", PePartitionPath)
		);
		if (!bStatus) { RestorePartition(PePartitionLetter, ShrinkPartitionLetter); return false; }
		NFileCount++;
		ui.progressBar_Status->setValue(30 + (int)(((NFileCount / FileCount) * 50)));

		//映像文件
		if (!bStatus) { RestorePartition(PePartitionLetter, ShrinkPartitionLetter); return false; }
		NFileCount++;
		ui.progressBar_Status->setValue(30 + (int)(((NFileCount / FileCount) * 50)));

		//其他文件
		for (auto Info : ExtendCopyFileList) {
			bStatus = AstCopyFile(
				Info.Src,
				Info.Des
			);
			if (!bStatus) { RestorePartition(PePartitionLetter, ShrinkPartitionLetter); return false; }

			NFileCount++;
			ui.progressBar_Status->setValue(30 + (int)(((NFileCount / FileCount) * 50)));
		}
	}
	catch (std::wstring e) {
		AstThrowError_ThrowError(
			ERROR_RUNTIME_CREATEPE_COPYFILE_FAIL,
			e
		);
		RestorePartition(PePartitionLetter, ShrinkPartitionLetter);
		return false;
	}
	AstLog_WriteMessage(
		LogHandle,
		L"[RUNTIME][CREATEPE][COPYFILE]文件复制完成！"
	);

	/* 创建Bcd启动项*/
	AstLog_WriteMessage(
		LogHandle,
		L"[RUNTIME][CREATEPE][BCDBOOT]开始创建BCD启动项！"
	);

	bStatus = CreateStartupEntry(_wstr2str(PePartitionPath)[0]);

	//创建失败，报错返回
	if (!bStatus) {
		AstThrowError_ThrowError(
			ERROR_RUNTIME_CREATEPE_BCDBOOT_CREATE_STARTUP_ENTRY_FAIL,
			L"创建启动项失败！"
		);
		RestorePartition(PePartitionLetter, ShrinkPartitionLetter);
		return false;
	}
	AstLog_WriteMessage(
		LogHandle,
		L"[RUNTIME][CREATEPE][BCDBOOT]BCD启动项创建完成！"
	);

	ui.progressBar_Status->setValue(95);

	/* 配置文件写入 */
	AstLog_WriteMessage(
		LogHandle,
		L"[RUNTIME][CREATEPE][CONFIG]开始写入配置文件"
	);

	try {
		//创建ini文件
		std::ifstream f(
			std::format(L"{}config.ini", PePartitionPath),
			std::fstream::out
		);
		f.close();

		//开始写入配置文件
		inicpp::IniManager _ConfigIni(
			std::format("{}config.ini", _wstr2str(PePartitionPath))
		);
		_ConfigIni.modify("config", "InstallFile", _wstr2str(std::format(L"Install.{}",ImageExtendName)));		//映像
		_ConfigIni.modify("config", "Index", SelectedIndex);	//序号
		//_ConfigIni.modify("config", "CopyLogToDesktop", CopyLogToDesktop);

	}
	catch (std::wstring e) {
		AstThrowError_ThrowError(
			ERROR_RUNTIME_CREATEPE_BCDBOOT_CREATE_STARTUP_ENTRY_FAIL,
			std::format(L"创建配置文件失败：\n{}", e)
		);
		RestorePartition(PePartitionLetter, ShrinkPartitionLetter);
		return false;
	}
	AstLog_WriteMessage(
		LogHandle,
		L"[RUNTIME][CREATEPE][CONFIG]配置文件写入成功！"
	);

	ui.progressBar_Status->setValue(99);

	/* 创建ID文件 */
	AstLog_WriteMessage(
		LogHandle,
		L"[RUNTIME][CREATEPE][CRETEID]开始创建识别文件！"
	);
	//先清掉各种残留
	for (auto iter : AstGetUsedLetter()) {
		std::wstring RestorePartitionFilePath = std::format(L"{}:\\{}", iter, AstRestorePartitionIDW);
		std::wstring PePartitionFilePath = std::format(L"{}:\\{}", iter, AstPePartitionIDW);
		std::wstring EfiPartitionFilePath = std::format(L"{}:\\{}", iter, AstEfiPartitionIDW);
		std::wstring ShrinkPartitionFilePath = std::format(L"{}:\\{}", iter, AstShrinkPartitionIDW);
		if (boost::filesystem::exists(RestorePartitionFilePath)) {
			AstLog_WriteWarning(
				LogHandle,
				std::format(L"[RUNTIME][CREATEPE][CRETEID]存在残留识别文件“{}”，正在删除", RestorePartitionFilePath)
			);
			remove(_wstr2str(RestorePartitionFilePath).c_str());
		}

		if (boost::filesystem::exists(PePartitionFilePath)) {
			AstLog_WriteWarning(
				LogHandle,
				std::format(L"[RUNTIME][CREATEPE][CRETEID]存在残留识别文件“{}”，正在删除", PePartitionFilePath)
			);
			remove(_wstr2str(PePartitionFilePath).c_str());
		}

		if (boost::filesystem::exists(EfiPartitionFilePath)) {
			AstLog_WriteWarning(
				LogHandle,
				std::format(L"[RUNTIME][CREATEPE][CRETEID]存在残留识别文件“{}”，正在删除", EfiPartitionFilePath)
			);
			remove(_wstr2str(EfiPartitionFilePath).c_str());
		}

		if (boost::filesystem::exists(ShrinkPartitionFilePath)) {
			AstLog_WriteWarning(
				LogHandle,
				std::format(L"[RUNTIME][CREATEPE][CRETEID]存在残留识别文件“{}”，正在删除", ShrinkPartitionFilePath)
			);
			remove(_wstr2str(ShrinkPartitionFilePath).c_str());
		}
	}

	try {
		std::wstring RestorePartitionFilePath = std::format(L"{}:\\{}", SystemReinstallPartition, AstRestorePartitionIDW);
		std::wstring PePartitionFilePath = std::format(L"{}:\\{}", PePartitionLetter, AstPePartitionIDW);
		std::wstring EfiPartitionFilePath = std::format(L"{}:\\{}", EfiPartitionLetter, AstEfiPartitionIDW);

		//创建文件
		AstLog_WriteMessage(
			LogHandle,
			std::format(L"[RUNTIME][CREATEPE][CRETEID]正在创建文件“{}”", RestorePartitionFilePath)
		);
		std::ofstream outfile;
		outfile.open(RestorePartitionFilePath);
		outfile.close();

		AstLog_WriteMessage(
			LogHandle,
			std::format(L"[RUNTIME][CREATEPE][CRETEID]正在创建文件“{}”", PePartitionFilePath)
		);
		outfile.open(PePartitionFilePath);
		outfile.close();

		AstLog_WriteMessage(
			LogHandle,
			std::format(L"[RUNTIME][CREATEPE][CRETEID]正在创建文件“{}”", EfiPartitionFilePath)
		);
		outfile.open(EfiPartitionFilePath);
		outfile.close();

		if (ShrinkPartitionLetter != NULL) {
			std::wstring ShrinkPartitionFilePath = std::format(L"{}:\\{}", ShrinkPartitionLetter, AstShrinkPartitionIDW);

			AstLog_WriteMessage(
				LogHandle,
				std::format(L"[RUNTIME][CREATEPE][CRETEID]正在创建文件“{}”", ShrinkPartitionFilePath)
			);
			outfile.open(ShrinkPartitionFilePath);
			outfile.close();
		}
		else {
			AstLog_WriteMessage(
				LogHandle,
				L"[RUNTIME][CREATEPE][CRETEID]未创建压缩分区识别文件，因为没有分区被压缩。"
			);
		}
	}
	catch (std::wstring e) {
		AstThrowError_ThrowError(
			ERROR_RUNTIME_CREATEPE_CREATEID_CREATEID_FAIL,
			e
		);
		RestorePartition(PePartitionLetter, ShrinkPartitionLetter);
		return false;
	}

	ui.progressBar_Status->setValue(100);

	AstLog_WriteMessage(
		LogHandle,
		L"[RUNTIME][CREATEPE]PE分区创建完成！"
	);

	return true;
}

void RestorePartition(
	WCHAR CreatedPartition,
	WCHAR ShrinkedPartition
) {
	CHAR CreatedPartitionA = _wstr2str(*new std::wstring{ CreatedPartition })[0];
	CHAR ShrinkedPartitionA = _wstr2str(*new std::wstring{ ShrinkedPartition })[0];
	AstDiskPart_DeleteVolume(CreatedPartitionA);
	if (ShrinkedPartition != NULL) {
		AstDiskPart_ExtendVolume(ShrinkedPartitionA);
	}
}

// 通过打开文件 获得文件大小，单位字节
long long getFileSize2(const char* fileName) {
	/*
		解释这里为什么用_fseeki64而不用fseek：
			_fseeki64可以兼容x64项目和x86项目
			fseek只能兼容x86项目

		这里说的兼容并不是能不能用的问题，而是准不准确的问题！
		在x64项目使用fseek，返回的数据可能会不准确！
	*/

	if (fileName == NULL) {
		//return -1;

		throw L"文件名不能为空！";
	}

	FILE* fp = NULL;
	long long fileSize = 0;

	// 打开文件
	errno_t err = fopen_s(&fp, fileName, "r");
	if (err != 0) {	// 不等0，打开文件失败
		throw std::format(L"打开文件“{}”失败！", _str2wstr(fileName));
	}

	// 将文件光标位置从末尾处移动0个位置 -->也可以这样说：(将文件光标从开头位置移动到末尾位置)
	int errer = _fseeki64(fp, 0, SEEK_END);
	if (err != 0) {	// 不等0，移动文件光标失败
		return -1;
	}

	// 获得当前文件指针位置，也就是间接获得文件大小
	fileSize = _ftelli64(fp);

	/* 如果只是单纯的获取文件大小，后续没有对文件有其他操作，下面将文件指针移动回来的代码可以忽略 */
	// 将文件光标位置从开头处移动0个位置 -->也可以这样说：(将文件光标从末尾位置移动到开头位置)
	errer = _fseeki64(fp, 0, SEEK_SET);
	if (err != 0) {	// 不等0，移动文件光标失败
		throw std::format(L"获取文件“{}”大小失败！", _str2wstr(fileName));
	}

	// 关闭文件
	fclose(fp);

	return fileSize;
}

/// <summary>
/// 复制文件
/// </summary>
/// <param name="Src">源文件</param>
/// <param name="Dst">新文件</param>
/// <returns></returns>
bool AstCopyFile(
	std::wstring Src,
	std::wstring Dst
) {
	BOOL bStatus = CopyFile(Src.c_str(), Dst.c_str(), FALSE);
	if (bStatus == 0) {
		AstThrowError_ThrowError(
			ERROR_RUNTIME_CREATEPE_COPYFILE_FAIL,
			std::format(L"[RUNTIME][CREATEPE][COPYFILE]将“{}”拷贝到“{}”发生错误，WinError代码：{}", Src, Dst, GetLastError())
		);
		return false;
	}
	AstLog_WriteMessage(
		LogHandle,
		std::format(L"[RUNTIME][CREATEPE][COPYFILE] {} -> {}", Src, Dst)
	);
	return true;
}

std::wstring AstGetFileExtendName(
	std::wstring filePath
) {
	return filePath.substr(filePath.find_last_of(L'.') + 1);
}


std::vector<WCHAR> AstGetUsedLetter() {
	std::vector<WCHAR> DriveLetter;

	//创建完全盘符列表
	std::vector<WCHAR> AllLetter;
	for (WCHAR wc = L'C'; wc <= L'Z'; wc++) {
		std::wstringstream LetterPath;
		LetterPath << L"\\\\.\\" << wc << L":";

		//尝试打开磁盘
		HANDLE hDevice = CreateFile(
			LetterPath.str().c_str(),
			GENERIC_READ,
			FILE_SHARE_READ | FILE_SHARE_WRITE,
			NULL,
			OPEN_EXISTING,
			0,
			NULL
		);

		if (hDevice != INVALID_HANDLE_VALUE) {
			CloseHandle(hDevice);
			DriveLetter.push_back(wc);
		}

	}

	if (DriveLetter.size() == 0) {
		return (std::vector<WCHAR>)NULL;
	}

	SetLastError(0);
	return DriveLetter;
}